
/*
 ---------------------------------------------------------------------------
 Copyright (c) 1998-2006, Brian Gladman, Worcester, UK. All rights reserved.

 LICENSE TERMS

 The free distribution and use of this software in both source and binary
 form is allowed (with or without changes) provided that:

   1. distributions of this source code include the above copyright
      notice, this list of conditions and the following disclaimer;

   2. distributions in binary form include the above copyright
      notice, this list of conditions and the following disclaimer
      in the documentation and/or other associated materials;

   3. the copyright holder's name is not used to endorse products
      built using this software without specific written permission.

 DISCLAIMER

 This software is provided 'as is' with no explicit or implied warranties
 in respect of its properties, including, but not limited to, correctness
 and/or fitness for purpose.
 ---------------------------------------------------------------------------
 Issue Date: 20/12/2007

 My thanks to Colin Sinclair for finding an error and suggesting a number
 of improvements to this code
*/

#include <stdio.h>
#include <math.h>
#include "../../common/types.h"
#include "inc/ctr.h"
#include "inc/mode_hdr.h"
#include "inc/debug.h"


#if defined(__cplusplus)
extern "C"
{
#endif

#define BLOCK_SIZE      AES_BLOCK_SIZE      /* block length                */
#define BLK_ADR_MASK    (BLOCK_SIZE - 1)    /* mask for 'in block' address */
#define INIT_STATE      10
#define AUTH_STATE      11
#define MSG_STATE       12

#define MAX_NONCE_LEN   13

#define ctr_len(x)      (((*(u8*)(x)) & 0x07) + 1)

#define clr_ctr(x,l)    memset((u8*)(x) + BLOCK_SIZE - (l), 0, (l))

#define set_ctr(x,v)                                \
  {                                                 \
    u8 *_p = (u8*)(x) + BLOCK_SIZE;       \
    length_t _t = (v), _l = ctr_len(x);             \
    do                                              \
    {                                               \
      *--_p = (u8)_t; _t >>= 8;                \
    } while (--_l);                                 \
  }

#if 0
#define inc_ctr(x)                                  \
  {                                                 \
    u8 *_p = (u8*)(x) + BLOCK_SIZE;       \
    length_t _l = ctr_len(x);                       \
    while (_l-- && ++(*--_p) == 0) ;                \
  }
#else
#define inc_ctr(x, m)                               \
  do {                                              \
      u8 *_p = (u8*)(x) + BLOCK_SIZE;     \
      if ((++(*--_p)) == ctx->mod)                  \
            *_p = 0;                                \
  }while(0);
#endif

#ifdef LONG_MESSAGES

#  define mlen_len(x)                     \
    (((x) & 0xff00000000000000ll) ? 7 :     \
     ((x) & 0xffff000000000000ll) ? 6 :     \
     ((x) & 0xffffff0000000000ll) ? 5 :     \
     ((x) & 0xffffffff00000000ll) ? 4 :     \
     ((x) & 0xffffffffff000000ll) ? 3 :     \
     ((x) & 0xffffffffffff0000ll) ? 2 : 1)

#else
#  define mlen_len(x) (((x) & 0xff000000) ? 3 : ((x) & 0xffff0000) ? 2 : 1)
#endif
extern int j;


void mod_inc_ctr(ctr_buf_t x, ctr_buf_t modVal, u8 mod)
{
    u8 i;

    if (!memcmp(UI8_PTR(x), UI8_PTR(modVal), BLOCK_SIZE)) {
        /* identical, that means counter go to the UPP value, reload to 0 */
        if (mod >=0 && mod <= 15) {
            memset(UI8_PTR(x), 0, BLOCK_SIZE);
        }
        else {
            u8 byteNum = mod / 8;
            u8 bitNum  = mod % 8;
            for (i = 0; i < byteNum; i++) {
                UI8_PTR(x)[BLOCK_SIZE - i - 1] = 0;
            }

            for(i = 0; i < bitNum; i++) {
                UI8_PTR(x)[BLOCK_SIZE - byteNum - 1] &= ~1<<i;
            }
        }


    }
    else {
        u8 *_p = (u8*)(x) + BLOCK_SIZE;
        do {
            ++(*--_p);
        } while(!*_p);
    }



}


/*****************************************************************************/

u32 ctr_init_and_enkey(                 /* initialise mode and set key */
            const unsigned char key[],     /* the key value               */
            unsigned long key_len,         /* and its length in bytes     */
            ctr_ctx ctx[1])                /* the mode context            */
{
  ctx->state = RETURN_GOOD;
  if ((key_len != 16) && (key_len != 24) && (key_len != 32))
  {
    ctx->state = CTR_BAD_KEY;
  }
  aes_encrypt_key(key, (int)key_len, ctx->aes);
  return RETURN_GOOD;
}

u32 ctr_init_and_dekey(                 /* initialise mode and set key */
            const unsigned char key[],     /* the key value               */
            unsigned long key_len,         /* and its length in bytes     */
            ctr_ctx ctx[1])                /* the mode context            */
{
  ctx->state = RETURN_GOOD;
  if ((key_len != 16) && (key_len != 24) && (key_len != 32))
  {
    ctx->state = CTR_BAD_KEY;
  }
  aes_decrypt_key(key, (int)key_len, (aes_decrypt_ctx*)&ctx->aes);
  return RETURN_GOOD;
}

/*****************************************************************************/

u32 ctr_init_message(                 /* initialise for a new message   */
            const unsigned char iv[],      /* the initialisation vector      */
            unsigned long iv_len,          /* the nonce length               */
            u32 initCounter,          /* the initialisation counter     */
            u8  mod,                  /* the initialisation counter mod */
            ctr_ctx ctx[1])                /* the mode context               */
{
  u32 cnt;
  u8  nonceLen;
  u8  byteNum;
  u8  bitNum;
  u8  i;

  ctx->state = RETURN_GOOD;

  ctx->mod = mod;
  ctx->cnt = 0;
  ctx->txt_acnt = 0;


#if 1
  memset(UI8_PTR(ctx->ctr_val), 0, 16);

  for(cnt = 0;
      cnt < iv_len;
      ++cnt)
  {
    UI8_PTR(ctx->ctr_val)[cnt] = iv[cnt];
  }

  /* initialize T1 */
  memcpy(UI8_PTR(ctx->mod_val),UI8_PTR(ctx->ctr_val), BLOCK_SIZE);

  if (ctx->mod >= 0 && ctx->mod <= 15) {
      memset(UI8_PTR(ctx->mod_val), 0xff, BLOCK_SIZE);
  }
  else {
      byteNum = ctx->mod / 8;
      bitNum  = ctx->mod % 8;

      for (i = 0; i < byteNum; i++) {
          UI8_PTR(ctx->mod_val)[BLOCK_SIZE - i - 1] |= 0xff;
      }

      for (i = 0; i < bitNum; i++) {
          UI8_PTR(ctx->mod_val)[BLOCK_SIZE-byteNum - 1] |= 1<<i;
      }
  }


#else

  UI8_PTR(ctx->ctr_val)[0] = 0x01;

  /* move the iv into the block */
  for(cnt = 1;
      cnt < (u32)BLOCK_SIZE - UI8_PTR(ctx->ctr_val)[0] - 1;
      ++cnt)
  {
    UI8_PTR(ctx->ctr_val)[cnt] = iv[cnt - 1];
  }

  /* set initial counter     */
  UI8_PTR(ctx->ctr_val)[BLOCK_SIZE - 1] = 1;


#endif


  return RETURN_GOOD;
}

/*****************************************************************************/

u32 ctr_encrypt_data(                   /* encrypt or decrypt data */
            unsigned char data[],          /* the data buffer         */
            unsigned long data_len,        /* and its length in bytes */
            ctr_ctx ctx[1])                /* the mode context        */
{
  u32  cnt = 0;
  u32  b_pos = ((u32)ctx->cnt) & BLK_ADR_MASK;


  if (ctx->state < 0)
  {
    //dbgp_ccm("   ctx->state\n");
    return ctx->state;
  }

  if (!data_len)
  {
    //dbgp_ccm("   ctx->data_len = 0\n");
    return RETURN_GOOD;
  }

  if (b_pos == 0)
  {
    aes_encrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), ctx->aes);
    //inc_ctr(ctx->ctr_val, ctx->mod);
    mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
  }

  if (!((data - (UI8_PTR(ctx->enc_ctr) + b_pos)) & BUF_ADRMASK))
  {
    while ((cnt < data_len) && (b_pos & BUF_ADRMASK))
    {
      data[cnt++] ^= UI8_PTR(ctx->enc_ctr)[b_pos++];
    }

    while ((cnt + BUF_INC <= data_len) && (b_pos <= BLOCK_SIZE - BUF_INC))
    {
      *UNIT_PTR(data + cnt) ^= *UNIT_PTR(UI8_PTR(ctx->enc_ctr) + b_pos);
      cnt += BUF_INC; b_pos += BUF_INC;
    }

    while (cnt + BLOCK_SIZE <= data_len)
    {
      aes_encrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), ctx->aes);
      mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
      xor_block_aligned(data + cnt, data + cnt, ctx->enc_ctr);
      cnt += BLOCK_SIZE;
    }
  }
  else
  {
    while ((cnt < data_len) && (b_pos < BLOCK_SIZE))
    {
      data[cnt++] ^= UI8_PTR(ctx->enc_ctr)[b_pos++];
    }

    while (cnt + BLOCK_SIZE <= data_len)
    {
      aes_encrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), ctx->aes);
      mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
      xor_block(data + cnt, data + cnt, ctx->enc_ctr);
      cnt += BLOCK_SIZE;
    }
  }

  while (cnt < data_len)
  {
    if (b_pos == BLOCK_SIZE)
    {
      aes_encrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), ctx->aes);
      mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
      b_pos = 0;
    }
    data[cnt++] ^= UI8_PTR(ctx->enc_ctr)[b_pos++];
  }

  ctx->cnt += cnt;
  return RETURN_GOOD;
}


u32 ctr_decrypt_data(                   /* encrypt or decrypt data */
            unsigned char data[],          /* the data buffer         */
            unsigned long data_len,        /* and its length in bytes */
            ctr_ctx ctx[1])                /* the mode context        */
{
  u32  cnt = 0;
  u32  b_pos = ((u32)ctx->cnt) & BLK_ADR_MASK;


  if (ctx->state < 0)
  {
    //dbgp_ccm("   ctx->state\n");
    return ctx->state;
  }

  if (!data_len)
  {
    //dbgp_ccm("   ctx->data_len = 0\n");
    return RETURN_GOOD;
  }

  if (b_pos == 0)
  {
    aes_decrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), (aes_decrypt_ctx*)&ctx->aes);
    //inc_ctr(ctx->ctr_val, ctx->mod);
    mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
  }

  if (!((data - (UI8_PTR(ctx->enc_ctr) + b_pos)) & BUF_ADRMASK))
  {
    while ((cnt < data_len) && (b_pos & BUF_ADRMASK))
    {
      data[cnt++] ^= UI8_PTR(ctx->enc_ctr)[b_pos++];
    }

    while ((cnt + BUF_INC <= data_len) && (b_pos <= BLOCK_SIZE - BUF_INC))
    {
      *UNIT_PTR(data + cnt) ^= *UNIT_PTR(UI8_PTR(ctx->enc_ctr) + b_pos);
      cnt += BUF_INC; b_pos += BUF_INC;
    }

    while (cnt + BLOCK_SIZE <= data_len)
    {
      aes_decrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), (aes_decrypt_ctx*)&ctx->aes);
      mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
      xor_block_aligned(data + cnt, data + cnt, ctx->enc_ctr);
      cnt += BLOCK_SIZE;
    }
  }
  else
  {
    while ((cnt < data_len) && (b_pos < BLOCK_SIZE))
    {
      data[cnt++] ^= UI8_PTR(ctx->enc_ctr)[b_pos++];
    }

    while (cnt + BLOCK_SIZE <= data_len)
    {
      aes_decrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), (aes_decrypt_ctx*)&ctx->aes);
      mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
      xor_block(data + cnt, data + cnt, ctx->enc_ctr);
      cnt += BLOCK_SIZE;
    }
  }

  while (cnt < data_len)
  {
    if (b_pos == BLOCK_SIZE)
    {
      aes_decrypt(UI8_PTR(ctx->ctr_val), UI8_PTR(ctx->enc_ctr), (aes_decrypt_ctx*)&ctx->aes);
      mod_inc_ctr(ctx->ctr_val, ctx->mod_val, ctx->mod);
      b_pos = 0;
    }
    data[cnt++] ^= UI8_PTR(ctx->enc_ctr)[b_pos++];
  }

  ctx->cnt += cnt;
  return RETURN_GOOD;
}

/*****************************************************************************/

u32 ctr_encrypt_message(              /* encrypt an entire message */
            const unsigned char iv[],      /* the initialisation vector */
            unsigned long iv_len,          /* and its length in bytes   */
            u32 mod,                  /* counter mod */
            u32 initCounter,          /* initialize counter */
            length_t msg_len,              /* and its length in bytes   */
            unsigned char msg[],           /* the message buffer        */
            ctr_ctx ctx[1])                /* the mode context          */
{

    u32 retVal;

    do {

        if (msg_len == 0) {
        }

        retVal = ctr_init_message(iv, iv_len, initCounter, mod, ctx);
        if (RETURN_GOOD != retVal) {
            break;
        }

        retVal = ctr_encrypt_data(msg, msg_len, ctx);
        if (retVal != RETURN_GOOD) {
            break;
        }


    }while(0);

    return retVal;

}



/*****************************************************************************/

u32 ctr_decrypt_message(              /* encrypt an entire message */
            const unsigned char iv[],      /* the initialisation vector */
            unsigned long iv_len,          /* and its length in bytes   */
            u32 mod,                  /* counter mod */
            u32 initCounter,          /* initialize counter */
            length_t msg_len,              /* and its length in bytes   */
            unsigned char msg[],           /* the message buffer        */
            ctr_ctx ctx[1])                /* the mode context          */
{

    u32 retVal;

    do {

        retVal = ctr_init_message(iv, iv_len, initCounter, mod, ctx);
        if (RETURN_GOOD != retVal) {
            break;
        }

        retVal = ctr_encrypt_data(msg, msg_len, ctx);
        if (retVal != RETURN_GOOD) {
            break;
        }

    }while(0);

    return retVal;

}





#if defined(__cplusplus)
}
#endif
